//以当前位置做力控
#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <tf2_ros/transform_listener.h>
#include <tf2_ros/buffer.h>
#include <eigen3/Eigen/Dense>
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>

using namespace std::chrono_literals;

static const rclcpp::Logger LOGGER = rclcpp::get_logger("force_servo_control_node");

class ForceServoControlNode : public rclcpp::Node
{
public:
    ForceServoControlNode() : Node("force_servo_control_node"), count_(0)
    {
        // 力控参数初始化
        // M_position_ = Eigen::Matrix3d::Identity() * 0.3;
        // K_position_ = Eigen::Matrix3d::Identity() * 100.0;
        // B_position_ = Eigen::Matrix3d::Identity() * 10.0;
        // // 在构造函数中修改这些参数
        M_position_ = Eigen::Matrix3d::Identity() * 1.0;  // 机械臂响应太慢（感觉"笨重"）→ 减小 M
        K_position_ = Eigen::Matrix3d::Identity() * 18.0; // 机械臂末端对位置偏差的恢复力
        B_position_ = Eigen::Matrix3d::Identity() * 13.0;  // 越小系统运动更"自由"，但可能产生振荡
        force_torque_reference_ = Eigen::Vector3d(0.0, 0.0, 0.0);
        velocity_ = Eigen::Vector3d::Zero();
        dt_ = 0.01;
        current_force_ = Eigen::Vector3d::Zero();

        // 创建力传感器数据订阅者
        force_sub_ = create_subscription<geometry_msgs::msg::TwistStamped>(
            "/force_data", 10, std::bind(&ForceServoControlNode::force_callback, this, std::placeholders::_1));

        // 初始化TF监听器
        tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
        tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);

        // 设置发布者
        twist_cmd_pub_ = create_publisher<geometry_msgs::msg::TwistStamped>("force_servo_control_node/delta_twist_cmds", 10);

        // 注意：不在构造函数中调用 setupPlanningSceneMonitor 和 setupServo
    }

    // 初始化函数，在对象构造完成后调用
    void initialize()
    {
        setupPlanningSceneMonitor();
        setupServo();

        // 获取机械臂的当前位姿作为参考位置
        getCurrentPositionAsReference();

        // 创建定时器，25Hz运行控制循环（40ms）
        timer_ = create_wall_timer(std::chrono::milliseconds(40), std::bind(&ForceServoControlNode::control_loop, this));

        RCLCPP_INFO(this->get_logger(), "力控伺服节点已初始化");
    }

private:
    rclcpp::Subscription<geometry_msgs::msg::TwistStamped>::SharedPtr force_sub_;
    rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
    std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
    std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
    std::shared_ptr<planning_scene_monitor::PlanningSceneMonitor> planning_scene_monitor_;
    std::unique_ptr<moveit_servo::Servo> servo_;
    rclcpp::TimerBase::SharedPtr timer_;

    Eigen::Vector3d current_force_;
    Eigen::Vector3d current_position_;
    Eigen::Vector3d velocity_;
    Eigen::Matrix3d M_position_;
    Eigen::Matrix3d K_position_;
    Eigen::Matrix3d B_position_;
    Eigen::Vector3d force_torque_reference_;
    Eigen::Vector3d position_reference_;  // 存储参考位置
    double dt_;
    size_t count_;

    // 获取机械臂当前的位姿并设置为参考位置
    void getCurrentPositionAsReference()
    {
        geometry_msgs::msg::TransformStamped transform;
        try {
            transform = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero);
            position_reference_.x() = transform.transform.translation.x;
            position_reference_.y() = transform.transform.translation.y;
            position_reference_.z() = transform.transform.translation.z;
            RCLCPP_INFO(this->get_logger(), "设置参考位置为: x=%.3f, y=%.3f, z=%.3f",
                        position_reference_.x(), position_reference_.y(), position_reference_.z());
        } catch (tf2::TransformException &ex) {
            RCLCPP_WARN(this->get_logger(), "无法获取TF变换: %s", ex.what());
        }
    }

    void force_callback(const geometry_msgs::msg::TwistStamped::SharedPtr msg)
    {
        current_force_.x() = -msg->twist.linear.z;
        current_force_.y() = 0.0;
        current_force_.z() = 0.0;
        // current_force_.z() = msg->twist.linear.x;
        // RCLCPP_INFO(this->get_logger(), "接收到力数据: x=%.3f, y=%.3f, z=%.3f",
        //             current_force_.x(), current_force_.y(), current_force_.z());
    }

    void control_loop()
    {
        geometry_msgs::msg::TransformStamped transform;
        try {
            transform = tf_buffer_->lookupTransform("base_link", "end", tf2::TimePointZero);
            current_position_.x() = transform.transform.translation.x;
            current_position_.y() = transform.transform.translation.y;
            current_position_.z() = transform.transform.translation.z;
        } catch (tf2::TransformException &ex) {
            RCLCPP_WARN(this->get_logger(), "无法获取TF变换: %s", ex.what());
            return;
        }
        Eigen::Vector3d force_difference = current_force_ - force_torque_reference_;
        Eigen::Vector3d position_difference = current_position_ - position_reference_;
        
        Eigen::Vector3d kp_position_term = K_position_ * position_difference;
        Eigen::Vector3d b_velocity_term = B_position_ * velocity_;
        
        Eigen::Vector3d acceleration = M_position_.inverse() * (force_difference - kp_position_term - b_velocity_term);
        velocity_ += acceleration * dt_;

        auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
        msg->header.stamp = this->now();
        msg->header.frame_id = "base_link";
        
        msg->twist.linear.x = velocity_.x();
        msg->twist.linear.y = velocity_.y();
        msg->twist.linear.z = velocity_.z();
        msg->twist.angular.x = 0.0;
        msg->twist.angular.y = 0.0;
        msg->twist.angular.z = 0.0;

        twist_cmd_pub_->publish(std::move(msg));

        // RCLCPP_INFO(this->get_logger(), "当前速度: x=%.3f, y=%.3f, z=%.3f",
        //             velocity_.x(), velocity_.y(), velocity_.z());
        count_++;
    }

    void setupPlanningSceneMonitor()
    {
        planning_scene_monitor_ = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
            shared_from_this(), "robot_description", tf_buffer_, "planning_scene_monitor");

        if (planning_scene_monitor_->getPlanningScene()) {
            planning_scene_monitor_->startStateMonitor("/joint_states");
            planning_scene_monitor_->setPlanningScenePublishingFrequency(25);
            planning_scene_monitor_->startPublishingPlanningScene(
                planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
                "/moveit_servo/publish_planning_scene");
            planning_scene_monitor_->startSceneMonitor();
            planning_scene_monitor_->providePlanningSceneService();
        } else {
            RCLCPP_ERROR(LOGGER, "规划场景配置失败");
            throw std::runtime_error("Planning scene setup failed");
        }
    }

    void setupServo()
    {
        auto servo_parameters = moveit_servo::ServoParameters::makeServoParameters(shared_from_this());
        if (!servo_parameters) {
            RCLCPP_FATAL(LOGGER, "无法加载伺服参数");
            throw std::runtime_error("Failed to load servo parameters");
        }

        servo_ = std::make_unique<moveit_servo::Servo>(shared_from_this(), servo_parameters, planning_scene_monitor_);
        servo_->start();
    }
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<ForceServoControlNode>();
    node->initialize();  // 在对象构造完成后调用初始化
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}
